# 使用RNN模型预测正弦波的示例

import numpy as np
import matplotlib.pyplot as plt
from rnn_model import RNN


# 生成正弦波数据
def generate_sine_wave(samples=200, period=20, amplitude=1.0):
    x = np.linspace(0, samples / period * 2 * np.pi, samples)
    y = amplitude * np.sin(x)
    return y


# 准备训练数据
def prepare_data(data, seq_length=10):
    x_data = []
    y_data = []

    for i in range(len(data) - seq_length):
        x_seq = []
        for j in range(seq_length):
            x_seq.append(np.array([[data[i + j]]]))

        y_seq = []
        for j in range(seq_length):
            y_seq.append(np.array([[data[i + j + 1]]]))

        x_data.append(x_seq)
        y_data.append(y_seq)

    return x_data, y_data


# 生成数据
sine_data = generate_sine_wave(samples=200)

# 数据归一化
data_max = np.max(sine_data)
data_min = np.min(sine_data)
sine_data_normalized = (sine_data - data_min) / (data_max - data_min)

# 准备训练数据
seq_length = 10
x_data, y_data = prepare_data(sine_data_normalized, seq_length)

# 划分训练集和测试集
train_size = int(len(x_data) * 0.8)
x_train, x_test = x_data[:train_size], x_data[train_size:]
y_train, y_test = y_data[:train_size], y_data[train_size:]

# 创建RNN模型
input_dim = 1
hidden_dim = 16
output_dim = 1
rnn = RNN(input_dim, hidden_dim, output_dim, lr=0.01)

# 训练模型
loss_history = []
epochs = 100

for epoch in range(epochs):
    epoch_loss = 0
    for i in range(len(x_train)):
        # 初始化隐藏状态
        h = np.zeros((hidden_dim, 1))
        h_seq = [h]
        y_pred_seq = []

        # 前向传播
        for t in range(seq_length):
            h = rnn.forward(x_train[i][t], h)
            y_pred = rnn.output(h)

            h_seq.append(h)
            y_pred_seq.append(y_pred)

        # 反向传播
        loss = rnn.backward(x_train[i], y_train[i], y_pred_seq, h_seq[1:])
        epoch_loss += loss

    avg_loss = epoch_loss / len(x_train)
    loss_history.append(avg_loss)

    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {avg_loss}")

# 测试模型
predictions = []
for i in range(len(x_test)):
    pred_seq = rnn.predict(x_test[i])
    predictions.append(pred_seq[-1][0][0])  # 获取序列最后一个值的预测

# 将预测结果转换回原始范围
predictions = np.array(predictions) * (data_max - data_min) + data_min

# 获取真实值
actual_values = []
for i in range(len(y_test)):
    actual_values.append(y_test[i][-1][0][0] * (data_max - data_min) + data_min)

# 可视化结果
plt.figure(figsize=(12, 6))
plt.plot(actual_values, label='Actual')
plt.plot(predictions, label='Predicted')
plt.legend()
plt.title('RNN Sine Wave Prediction')
plt.savefig('sine_prediction.png')
plt.show()

# 绘制损失曲线
plt.figure(figsize=(10, 5))
plt.plot(loss_history)
plt.title('Training Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.savefig('training_loss.png')
plt.show()